""" 优化版OCR web服务 """
import json
import os
import psutil
import socket
import tempfile

from concurrent.futures import as_completed, ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from time import time
from bottle import request, BaseRequest, response, Bottle
from PPOCR_api import GetOcrApi
from RapidOCR_api import OcrAPI


@dataclass
class OCRTask:
    type: str
    filename: str
    data: bytes or str = None
    ocr_type: str = "RapidOCR"
    original_filename: str = None
    page_num: int = None
    error: str = None


class OCRService:
    """ocr web 服务 """

    def __init__(self):
        """初始化OCR服务"""
        self.app = Bottle()

        self.setup_router()
        self.MAX_WORKERS = min(psutil.cpu_count(), 8)
        self.GLOBAL_THREAD_POOL = ThreadPoolExecutor(max_workers=self.MAX_WORKERS)  # 添加全局线程池
        self.WEB_SUPPORTED_OCR_TYPES = ["PaddleOCR", "RapidOCR"]  # 存储可用的OCR引擎类型
        self.SERVER_DIR = Path(__file__).resolve().parent
        self.Rapid_OCR_PATH = str(self.SERVER_DIR.joinpath("OCRCore/RapidOCR/RapidOCR.exe"))
        self.Paddle_OCR_PATH = str(self.SERVER_DIR.joinpath("OCRCore/PaddleOCR/PaddleOCR.exe"))
        self.TEST_OCR_FILE_PATH = str(self.SERVER_DIR.joinpath("images/test_ocr.png"))
        self.ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.pdf'}
        self.IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg'}
        self.OCR_PROCESS_POOL = {}

    def setup_router(self):
        """ 路由注册中心"""
        self.app.route('/ping', callback=self.ping)
        self.app.route('/ocr_engines', callback=self.ocr_engines)
        self.app.route('/ocr_preheat', callback=self.ocr_preheat)
        self.app.post('/ocr', callback=self.ocr)
        self.app.post('/ocr_sse_optimized', callback=self.ocr_sse_optimized)
        self.app.post('/ocr_sse_file', callback=self.ocr_sse_file)

    def ocr_engines(self):
        """测试本地主机支持的ocr引擎"""
        supported_ocr_types = []
        for engine_type in self.WEB_SUPPORTED_OCR_TYPES:
            result = None
            ocr_instance = None
            try:
                if engine_type == "RapidOCR":
                    ocr_instance = OcrAPI(self.Rapid_OCR_PATH)
                elif engine_type == "PaddleOCR":
                    ocr_instance = GetOcrApi(self.Paddle_OCR_PATH)

                result = ocr_instance.run(self.TEST_OCR_FILE_PATH)
                if result and result.get("code") == 100:
                    supported_ocr_types.append(engine_type)
            except Exception as e:
                print(f"初始化{engine_type}失败: {e}")
            finally:
                """释放资源"""
                if ocr_instance:
                    ocr_instance.stop() if hasattr(ocr_instance, 'stop') else ocr_instance.exit()

        if supported_ocr_types:
            return json.dumps(
                {"code": 100, "msg": "初始化成功", "supported_ocr_types": supported_ocr_types})
        else:
            return json.dumps({"code": 404, "msg": "未找到可用的OCR引擎"})

    def ocr_preheat(self):
        """预热ocr服务"""
        response.content_type = 'text/event-stream'
        response.cache_control = 'no-cache'

        def preheat_stream():
            """
            预热ocr服务
            Returns:

            """
            yield f"data: {json.dumps({'status': 'detecting', 'message': '正在检测可用OCR引擎'})}\n\n"
            supported_engines = json.loads(self.ocr_engines()).get('supported_ocr_types', [])
            if not supported_engines:
                yield f"data: {json.dumps({'status': 'error', 'error': '没有可用的OCR引擎'})}\n\n"
                yield "event: end\ndata: {}\n\n"
                return
            # 清空ocr实例池
            self.OCR_PROCESS_POOL = {}
            for engine_type in supported_engines:
                res_dict = {'status': 'preheating', 'engine': engine_type,
                            'message': f'正在预热{engine_type}'}
                yield f"data: {json.dumps(res_dict)}\n\n"
                self.OCR_PROCESS_POOL[engine_type] = []
                for i in range(self.MAX_WORKERS):
                    try:
                        if engine_type == "RapidOCR":
                            instance = OcrAPI(self.Rapid_OCR_PATH)
                        else:  # PaddleOCR
                            instance = GetOcrApi(self.Paddle_OCR_PATH)
                        # 测试实例是否可用
                        test_result = instance.run(self.TEST_OCR_FILE_PATH)
                        if test_result and test_result.get("code") == 100:
                            self.OCR_PROCESS_POOL[engine_type].append(instance)
                            res_dict = {'status': 'progress', 'engine': engine_type,
                                        'worker': i + 1, 'total': self.MAX_WORKERS}
                            yield f"data: {json.dumps(res_dict)}\n\n"
                    except Exception as e:
                        res_dict = {'status': 'error', 'engine': engine_type, 'error': str(e)}
                        yield f"data: {json.dumps(res_dict)}\n\n"
            res_dict = {'status': 'completed', 'message': 'OCR预热完成',
                        'engines': supported_engines, 'workers_per_engine': self.MAX_WORKERS}
            yield f"data: {json.dumps(res_dict)}\n\n"
            yield "event: end\ndata: {}\n\n"

        return preheat_stream()

    def ping(self):
        """测试接口"""
        return json.dumps({"code": 200, "msg": "pong"})

    def init_ocr_pool(self):
        """没有经过预热的ocr 默认使用 一个RapidOCR实例进行使用"""
        instance = OcrAPI(self.Rapid_OCR_PATH)
        test_result = instance.run(self.TEST_OCR_FILE_PATH)
        if test_result and test_result.get("code") == 100:
            self.OCR_PROCESS_POOL["RapidOCR"] = [instance]

    def ocr_sse_optimized(self):
        """SSE版本的OCR接口，适合100-200张图片"""
        response.content_type = 'text/event-stream'
        response.cache_control = 'no-cache'

        def generate_sse():
            upload_files = request.files.getall('file')
            if not upload_files:
                yield f"data: {json.dumps({'code': 400, 'msg': '请上传文件'})}\n\n"
                yield "event: end\ndata: {}\n\n"
                return

            start_time = time()
            tasks = []
            completed = 0

            with tempfile.TemporaryDirectory() as tmp_dir:
                # 任务准备阶段
                for upload_file in upload_files:
                    if not allowed_file(upload_file.raw_filename, self.ALLOWED_EXTENSIONS):
                        res_dict = {'status': 'error', 'filename': upload_file.raw_filename,
                                    'error': f'不支持的文件类型: {get_file_extension(upload_file.raw_filename)}'}
                        yield f"data: {json.dumps(res_dict)}\n\n"
                        continue

                    file_extension = get_file_extension(upload_file.raw_filename)
                    ocr_type = request.query.get("ocr_type", "RapidOCR")

                    if file_extension in self.IMAGE_EXTENSIONS:
                        tasks.append(OCRTask(
                            type='image_memory',
                            filename=upload_file.raw_filename,
                            data=upload_file.file.read(),
                            ocr_type=ocr_type
                        ))
                    elif file_extension == '.pdf':
                        file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                        upload_file.save(file_path)
                        image_files = pdf_to_images(file_path, tmp_dir,
                                                    Path(file_path).stem)
                        for i, image_path in enumerate(image_files):
                            tasks.append(OCRTask(
                                type='image_file',
                                filename=f"{upload_file.raw_filename}_page_{i + 1}",
                                data=image_path,
                                ocr_type=ocr_type,
                                original_filename=upload_file.raw_filename,
                                page_num=i + 1
                            ))

                total_tasks = len(tasks)
                res_dict = {'status': 'processing', 'total_tasks': total_tasks,
                            'message': f'开始处理{total_tasks}个任务'}
                yield f"data: {json.dumps(res_dict)}\n\n"

                # 并行处理阶段
                futures = []
                futures_results_mapping = {}
                for task in tasks:
                    if task.type == 'image_memory':
                        future = self.GLOBAL_THREAD_POOL.submit(
                            partial(self.process_image_memory_task, ocr_core_type=task.ocr_type),
                            task.data
                        )
                        futures_results_mapping[future] = task
                        futures.append(future)
                    elif task.type == 'image_file':
                        future = self.GLOBAL_THREAD_POOL.submit(
                            partial(self.process_image_file_task, ocr_core_type=task.ocr_type),
                            task.data
                        )
                        futures.append(future)
                        futures_results_mapping[future] = task

                for future in as_completed(futures):
                    try:
                        result = future.result()
                        task = futures_results_mapping[future]
                        completed += 1
                        result.update({
                            'status': 'completed',
                            'progress': f'{completed}/{total_tasks}',
                            'filename': task.filename,
                            'original_filename': task.original_filename,
                        })
                        yield f"data: {json.dumps(result)}\n\n"
                    except Exception as e:
                        completed += 1
                        res_dict = {'status': 'error', 'error': str(e),
                                    'progress': f'{completed}/{total_tasks}'}
                        yield f"data: {json.dumps(res_dict)}\n\n"

            total_time = time() - start_time
            res_dict = {'status': 'all_completed', 'total_time': total_time,
                        'avg_time_per_task': total_time / total_tasks if total_tasks else 0,
                        'workers_used': self.MAX_WORKERS}
            yield f"data: {json.dumps(res_dict)}\n\n"
            yield "event: end\ndata: {}\n\n"

        return generate_sse()

    def ocr(self):
        """ocr接口 适合 100-200 张图片"""
        upload_files = request.files.getall('file')

        if not upload_files:
            return json.dumps({"code": 400, "msg": "请上传文件"})
        tasks = []

        start_time = time()
        with tempfile.TemporaryDirectory() as tmp_dir:
            # 任务准备阶段
            for upload_file in upload_files:
                if not allowed_file(upload_file.raw_filename, self.ALLOWED_EXTENSIONS):
                    tasks.append(OCRTask(
                        type='error',
                        filename=upload_file.raw_filename,
                        error=f"不支持的文件类型: {get_file_extension(upload_file.raw_filename)}"
                    ))
                    continue

                file_extension = get_file_extension(upload_file.raw_filename)
                ocr_type = request.query.get("ocr_type", "RapidOCR")

                if file_extension in self.IMAGE_EXTENSIONS:
                    tasks.append(OCRTask(
                        type='image_memory',
                        filename=upload_file.raw_filename,
                        data=upload_file.file.read(),
                        ocr_type=ocr_type
                    ))
                elif file_extension == '.pdf':
                    file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                    upload_file.save(file_path)
                    image_files = pdf_to_images(file_path, tmp_dir, Path(file_path).stem)
                    for i, image_path in enumerate(image_files):
                        tasks.append(OCRTask(
                            type='image_file',
                            filename=f"{upload_file.raw_filename}_page_{i + 1}",
                            data=image_path,
                            ocr_type=ocr_type,
                            original_filename=upload_file.raw_filename,
                            page_num=i + 1
                        ))
            results = []
            futures = []
            futures_results_mapping = {}
            for task in tasks:
                if task.type == 'image_memory':
                    future = self.GLOBAL_THREAD_POOL.submit(
                        partial(self.process_image_memory_task, ocr_core_type=task.ocr_type),
                        task.data
                    )
                    futures.append(future)
                    futures_results_mapping[future] = task
                elif task.type == 'image_file':
                    future = self.GLOBAL_THREAD_POOL.submit(
                        partial(self.process_image_file_task, ocr_core_type=task.ocr_type),
                        task.data
                    )
                    futures.append(future)
                    futures_results_mapping[future] = task

            for future in as_completed(futures):
                try:
                    result = future.result()
                    result.update({
                        'filename': task.filename,
                        'status': 'completed',
                        'original_filename': task.original_filename,
                    })
                    results.append(result)
                except Exception as e:
                    results.append({"error": str(e)})
            total_time = time() - start_time
            return json.dumps({
                "code": 200,
                "msg": "success",
                "data": results,
                "performance": {
                    "total_time": total_time,
                    "avg_time_per_task": total_time / len(tasks) if tasks else 0,
                    "workers_used": self.MAX_WORKERS
                }
            })

    def ocr_sse_file(self):
        """ 前端 传递一个 格式为 pdf 的 大文件"""
        response.content_type = 'text/event-stream'
        response.cache_control = 'no-cache'

        def generate_sse():
            pdf_abs_paths = request.json.get('pdf_paths')
            if not pdf_abs_paths:
                yield f"data: {json.dumps({'code': 400, 'msg': '请上传pdf文件路径'})}\n\n"
                yield "event: end\ndata: {}\n\n"
                return
            start_time = time()
            tasks = []
            completed = 0
            with tempfile.TemporaryDirectory() as tmp_dir:
                for pdf_path in pdf_abs_paths:
                    filename = os.path.basename(pdf_path)
                    if not allowed_file(filename, self.ALLOWED_EXTENSIONS):
                        res_dict = {'status': 'error', 'filename': filename,
                                    'error': f'不支持的文件类型: {get_file_extension(filename)}'}
                        yield f"data: {json.dumps(res_dict)}\n\n"
                        continue
                    ocr_type = request.query.get("ocr_type", "RapidOCR")
                    res_dict = {'status': 'processing', 'filename': filename,
                                'message': '正在转换PDF为图片'}
                    yield f"data: {json.dumps(res_dict)}\n\n"

                    images_files = pdf_to_images(pdf_path, tmp_dir, Path(pdf_path).stem)
                    for i, image_path in enumerate(images_files):
                        tasks.append(OCRTask(
                            type='image_file',
                            filename=f"{filename}_page_{i + 1}",
                            data=image_path,
                            ocr_type=ocr_type,
                            original_filename=filename,
                            page_num=i + 1
                        ))
                    total_tasks = len(tasks)
                    res_dict = {'status': 'processing', 'total_tasks': total_tasks,
                                'message': f'开始处理{total_tasks}个任务'}
                    yield f"data: {json.dumps(res_dict)}\n\n"

                    results = []
                    futures = []
                    futures_results_mapping = {}
                    for task in tasks:
                        if task.type == 'image_file':
                            future = self.GLOBAL_THREAD_POOL.submit(
                                partial(self.process_image_file_task, ocr_core_type=task.ocr_type),
                                task.data
                            )
                            futures.append(future)
                            futures_results_mapping[future] = task
                    for future in as_completed(futures):
                        try:
                            task = futures_results_mapping[future]
                            result = future.result()
                            completed += 1
                            result.update({
                                'status': 'completed',
                                'filename': task.filename,
                                'original_filename': task.original_filename,
                                'page_num': task.page_num,
                                'progress': f'{completed}/{total_tasks}'
                            })
                            yield f"data: {json.dumps(result)}\n\n"
                        except Exception as e:
                            completed += 1
                            res_dict = {'status': 'error', 'error': str(e),
                                        'progress': f'{completed}/{total_tasks}'}
                            yield f"data: {json.dumps(res_dict)}\n\n"

                total_time = time() - start_time
                res_dict = {'status': 'all_completed', 'total_time': total_time,
                            'avg_time_per_task': total_time / total_tasks if total_tasks else 0,
                            'workers_used': self.MAX_WORKERS}
                yield f"data: {json.dumps(res_dict)}\n\n"
                yield "event: end\ndata: {}\n\n"

        return generate_sse()

    def process_image_memory_task(self, file_bytes, ocr_core_type: str = "RapidOCR"):
        """优化版内存图片处理"""
        ocr_core = None
        try:
            ocr_core = self.get_ocr_instance(ocr_core_type)
            if ocr_core_type == "PaddleOCR":
                return ocr_core.runBytes(file_bytes)
            else:
                return ocr_core.runBytes(file_bytes)
        finally:
            if ocr_core:
                self.release_ocr_instance(ocr_core_type, ocr_core)

    def process_image_file_task(self, image_path, ocr_core_type: str = "RapidOCR"):
        """优化版文件图片处理"""
        ocr_core = None
        try:
            ocr_core = self.get_ocr_instance(ocr_core_type)
            if ocr_core_type == "PaddleOCR":
                return ocr_core.run(image_path)
            else:
                return ocr_core.run(image_path)
        finally:
            if ocr_core:
                self.release_ocr_instance(ocr_core_type, ocr_core)

    def get_ocr_instance(self, ocr_type: str):
        """获取OCR实例"""

        if not self.OCR_PROCESS_POOL.get(ocr_type):
            self.init_ocr_pool()
        if not self.OCR_PROCESS_POOL.get(ocr_type):
            """没有可用的ocr实例"""
            raise Exception(f"没有可用的{ocr_type}实例")

        return self.OCR_PROCESS_POOL[ocr_type].pop(0)

    def release_ocr_instance(self, ocr_type: str, instance):
        """循环利用OCR实例"""
        self.OCR_PROCESS_POOL[ocr_type].append(instance)

    def start(self, port: int = 9911, memory_size=300 * 1024 * 1024):
        """启动服务"""
        BaseRequest.MEMFILE_MAX = memory_size
        try:
            self.app.run(port=port)
        finally:

            self.GLOBAL_THREAD_POOL.shutdown(wait=True)
            """释放ocr 实例对象"""
            for instance in self.OCR_PROCESS_POOL.values():
                instance.stop() if hasattr(instance, 'stop') else instance.exit()



def is_port_in_use(port):
    """
    检查端口是否被占用
    :param port:
    :return:
    """
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) == 0

def get_file_extension(filename):
    """
    获取文件扩展名
    :param filename:
    :return:
    """
    return os.path.splitext(filename)[1].lower()


def allowed_file(filename, ALLOWED_EXTENSIONS):
    """
    判断文件是否为允许的类型
    :param filename:
    :return:
    """
    return get_file_extension(filename) in ALLOWED_EXTENSIONS


def pdf_to_images(pdf_path, output_dir, file_name):
    """
    pdf转图片
    :param pdf_path: PDF文件路径
    :param output_dir: 输出目录
    :param file_name: 文件名前缀
    :return: 图片文件路径列表
    """
    import fitz
    pdf_doc = fitz.open(pdf_path)
    file_names = []
    for page_index in range(pdf_doc.page_count):
        page = pdf_doc[page_index]
        zoom = (1.3, 1.3)
        mat = fitz.Matrix(zoom[0], zoom[1])
        pix = page.get_pixmap(matrix=mat, alpha=False)
        temp_file_name = file_name + str(page_index) + '.png'
        image_abs_path = os.path.join(output_dir, temp_file_name)
        pix.save(image_abs_path)
        file_names.append(image_abs_path)
    return file_names


if __name__ == "__main__":
    service = OCRService()
    service.start()
